// HookEngine.cpp : Defines the exported functions for the DLL application.
//

#include "stdafx.h"
#include "HookEngine.h"
#include "HookSettingEntry.h"
#include "libdasm.h"
#include <Tlhelp32.h>
#include <Psapi.h>

#pragma comment(lib, "libdasm.lib")
#pragma comment(lib, "strsafe.lib")

static const int X86_PAGE_SIZE = 4*1024;
static const int SHORT_CALL_INSTRUCTION_LENGTH = 5;

static const int SHORT_JMP_INSTRUCTION_LENGTH = 5;
static const int MAX_X86_INSTRUCTION_LENGTH = 14;

const DWORD X86_SHORT_CALL_OPCODE_LENGTH = 1;
const DWORD X86_SHORT_JMP_OPCODE_LENGTH = 1;


const DWORD X86_NOP_INSTRUCTION = 0x90;
const DWORD X86_SHORTCALL_OPCODE = 0xE8;
const DWORD X86_SHORTJMP_OPCODE = 0xE9;

static const int MAX_SAVED_PROLOG_LENGTH = (SHORT_CALL_INSTRUCTION_LENGTH-1+MAX_X86_INSTRUCTION_LENGTH);
static const int MAX_HOOK_RECORD_COUNT = 100;



void __stdcall Prolog(DWORD_PTR * pHookedAPIReturnAddress);
void __stdcall Epilog(DWORD_PTR * pdwHookedAPIReturnAddress);
void ProxyEpilog();
void ProxyProlog();

#pragma pack(1)
struct StackTrace
{
	DWORD m_dwCount;
	DWORD_PTR m_dwpStackFrames[1];
};

struct StackTraceList
{
	DWORD m_dwCount;
	StackTrace m_StackTraces[1]; 
};

struct ModuleInfo
{
	DWORD64 m_dwpStartAddress;
	DWORD m_dwSize;
	wchar_t m_szFullPath[MAX_PATH];
};

struct ModuleInfoList
{
	DWORD m_dwCount;
	ModuleInfo m_modules[1];
};
#pragma pack()


static StackTraceList * s_pStackTraceList;
static StackTrace * s_pCurrentStackTrace;

StackTrace * AllocateNewStackTrace(DWORD dwStackFrameCount)
{
	assert(s_pCurrentStackTrace!=NULL);
	StackTrace *pOldPointer = NULL, *pNewPointer = NULL;
	do
	{
		pOldPointer = s_pCurrentStackTrace;
		pNewPointer = (StackTrace *)((BYTE *)pOldPointer+sizeof(DWORD_PTR)*dwStackFrameCount+sizeof(DWORD));
	}while(pOldPointer!=InterlockedCompareExchangePointer((PVOID volatile*)&s_pCurrentStackTrace, pNewPointer, pOldPointer));

	InterlockedIncrement((volatile LONG *)&s_pStackTraceList->m_dwCount);
	assert(pOldPointer!=NULL);
	if(pOldPointer==NULL)
	{
		NullPointerException ExceptionObj("pOldPointer is NULL\n");
		throw ExceptionObj;
	}

	pOldPointer->m_dwCount = dwStackFrameCount;
	return pOldPointer;
}

DWORD GetStackTraceListSize()
{
	assert(s_pStackTraceList!=NULL);
	if(s_pStackTraceList==NULL)
	{
		NullPointerException ExceptionObj("s_pStackTraceList is NULL\n");
		throw ExceptionObj;
	}

	if(s_pStackTraceList->m_dwCount==0)
	{
		return 0;
	}

	DWORD dwSize = sizeof(DWORD);
	
	TRACEA("[GetStackTraceListSize()] -- s_pStackTraceList count is %d\n", s_pStackTraceList->m_dwCount);
	StackTrace * pCurrentStackTrace = s_pStackTraceList->m_StackTraces;
	for(DWORD dwIndex=0;dwIndex<s_pStackTraceList->m_dwCount;dwIndex++)
	{
		TRACEA("[GetStackTraceListSize()] -- StackFrame[%d] count is %d\n", dwIndex, pCurrentStackTrace->m_dwCount);
		DWORD dwCurrentSize = pCurrentStackTrace->m_dwCount*sizeof(DWORD_PTR)+sizeof(DWORD);
		dwSize+=dwCurrentSize;

		pCurrentStackTrace = (StackTrace *)((BYTE *)pCurrentStackTrace+dwCurrentSize);
	}

	assert(dwSize<=::X86_PAGE_SIZE*1024);
	return dwSize;
}


class HookRecord;
struct HookedAPIStackFrame
{
	DWORD_PTR m_dwpReturnAddress;
	HookRecord * m_pHookRecord;

	HookedAPIStackFrame(){}
	HookedAPIStackFrame(DWORD_PTR dwpReturnAddress, HookRecord * pHookRecord): 
		m_dwpReturnAddress(dwpReturnAddress), m_pHookRecord(pHookRecord)
	{
	}
};

class x86CodeThunk
{
private:
	const static BYTE X86_THUNK_INSTRUCTINO_LENGTH = 5;
	const static BYTE X86_THUNK_OPRAND_LENGTH = sizeof(DWORD_PTR);
	BYTE m_ThunkInstruction[X86_THUNK_INSTRUCTINO_LENGTH];
public:
	x86CodeThunk()
	{
		memset(m_ThunkInstruction, X86_NOP_INSTRUCTION, X86_THUNK_INSTRUCTINO_LENGTH);
	}

	x86CodeThunk(BYTE bOpcode)
	{
		memset(m_ThunkInstruction, X86_NOP_INSTRUCTION, X86_THUNK_INSTRUCTINO_LENGTH);
		this->SetOpcode(bOpcode);
	}

	void SetOpcode(BYTE bOpcode)
	{
		m_ThunkInstruction[0] = bOpcode;
	}	

	void SetOprandRelativeAddress(DWORD_PTR dwAbsoluteAddress, DWORD_PTR dwInstructionAddressAfterThunk)
	{
		DWORD_PTR dwRelativeAddress  = dwAbsoluteAddress - dwInstructionAddressAfterThunk;
		*(DWORD_PTR *)&(m_ThunkInstruction[sizeof(BYTE)]) = dwRelativeAddress;
	}

	LPVOID GetCodeThunkAddress()
	{
		LPVOID pCodeThunkAddress = &m_ThunkInstruction;
		assert( pCodeThunkAddress>0 );
		return pCodeThunkAddress;
	}

	DWORD GetCodeThunkLength()
	{
		return sizeof(m_ThunkInstruction);
	}
};


class ReferenceCountEntry
{
	DWORD_PTR m_dwpKey;
	int m_nCount;
public:
	ReferenceCountEntry():m_nCount(0), m_dwpKey(0){}

	void SetKey(DWORD_PTR dwpKey){m_dwpKey=dwpKey;}
	DWORD_PTR GetKey(){return m_dwpKey;}
	DWORD_PTR GetCount(){return m_nCount;}

	DWORD CheckIncrease()
	{
		DWORD dwNewCount;
		int old;
		do
		{
			old = m_nCount;
			if(old>0)
			{
				dwNewCount = old+1;
			}
			else
			{
				dwNewCount = old;
			}
		}while(old!=InterlockedCompareExchange((LONG volatile*)&m_nCount, dwNewCount, old));
		
		return dwNewCount;
	}

	DWORD Increase()
	{
		return ::InterlockedIncrement((LONG volatile*)&m_nCount);
	}

	DWORD Decrease()
	{
		return ::InterlockedDecrement((LONG volatile*)&m_nCount);
	}
};

class HookAPIReferenceCountList
{
private:
	DWORD m_dwHookAPICount;
	ReferenceCountEntry m_CountEntry[MAX_HOOK_RECORD_COUNT];
public:
	HookAPIReferenceCountList():m_dwHookAPICount(0){}

	ReferenceCountEntry * FindEntry(DWORD_PTR dwpKey)
	{
		ReferenceCountEntry * pCountEntry = NULL;
		for(DWORD dwIndex=0;dwIndex<m_dwHookAPICount;dwIndex++)
		{
			if(m_CountEntry[dwIndex].GetKey()==dwpKey)
			{
				pCountEntry = &(m_CountEntry[dwIndex]);
				break;
			}
		}

		assert(pCountEntry!=NULL);
		return pCountEntry;
	}

	ReferenceCountEntry * InitializeNewEntry(DWORD_PTR dwpKey)
	{
		ReferenceCountEntry * pCountEntry = NULL;
		for(DWORD dwIndex=0;dwIndex<m_dwHookAPICount;dwIndex++)
		{
			if(m_CountEntry[dwIndex].GetKey()==dwpKey)
			{
				pCountEntry = &(m_CountEntry[dwIndex]);
				pCountEntry->Increase();
				return pCountEntry;
			}
		}

		m_CountEntry[m_dwHookAPICount].SetKey(dwpKey);
		m_CountEntry[m_dwHookAPICount].Increase();
		m_dwHookAPICount++;
		return &(m_CountEntry[m_dwHookAPICount]);
	}
};
static auto_ptr<HookAPIReferenceCountList> s_pHookAPIReferenceCountList;


#pragma pack(1)
class HookRecord
{
public:
	class Prolog_Thunk
	{
		BYTE m_SavedAPIEntryPointLength;

		BYTE m_SavedAPIEntryPointBuffer[MAX_SAVED_PROLOG_LENGTH];
		x86CodeThunk m_JmpThunkToAPIReEntry;
	public:
		Prolog_Thunk()
		{
			memset(this, X86_NOP_INSTRUCTION, sizeof(*this));
			m_JmpThunkToAPIReEntry.SetOpcode(X86_SHORTJMP_OPCODE);
			m_SavedAPIEntryPointLength = 0;
		}

		void SaveAPIEntryPoint(PVOID pAPIEntryAddress, DWORD dwOverrwrittenAPIEntryLength)
		{
			memcpy_s(m_SavedAPIEntryPointBuffer, sizeof(m_SavedAPIEntryPointBuffer), pAPIEntryAddress, dwOverrwrittenAPIEntryLength);
		}

		void SetSavedAPIEntryPointLength(BYTE bSavedAPIEntryLength)
		{
			m_SavedAPIEntryPointLength = bSavedAPIEntryLength;
		}

		BYTE GetSavedAPIEntryPointLength()
		{
			return m_SavedAPIEntryPointLength;
		}

		BYTE GetSavedAPIEntryPointBufferLength()
		{
			return sizeof(BYTE)*MAX_SAVED_PROLOG_LENGTH;
		}

		void SetAPIReEntryRelativeAddress(DWORD_PTR dwAPIReEntryAbsoluteAddress, DWORD_PTR dwNextInstructionAddress)
		{
			m_JmpThunkToAPIReEntry.SetOprandRelativeAddress( dwAPIReEntryAbsoluteAddress, dwNextInstructionAddress );
		}

		DWORD_PTR GetSavedAPIEntryPointBufferAddress()
		{
			return (DWORD_PTR)( this->m_SavedAPIEntryPointBuffer );
		}

		DWORD_PTR GetNextAddressOfJmpThunkToAPIReEntry()
		{
			return (DWORD_PTR)m_JmpThunkToAPIReEntry.GetCodeThunkAddress() + m_JmpThunkToAPIReEntry.GetCodeThunkLength();
		}
	};

	class Epilog_Thunk
	{
		x86CodeThunk m_CallThunkToProxyEpilog;
	public:
		Epilog_Thunk(): m_CallThunkToProxyEpilog(X86_SHORTCALL_OPCODE)
		{
			DWORD_PTR dwNextAddressOfCallInstruction = (DWORD_PTR)( GetCallThunkToProxyEpilogAddress() ) + m_CallThunkToProxyEpilog.GetCodeThunkLength();
			m_CallThunkToProxyEpilog.SetOprandRelativeAddress( (DWORD_PTR)(ProxyEpilog), dwNextAddressOfCallInstruction );
		}

		DWORD_PTR GetCallThunkToProxyEpilogAddress()
		{
			return (DWORD_PTR)( m_CallThunkToProxyEpilog.GetCodeThunkAddress() );
		}
	};

	Prolog_Thunk m_PrologThunk;
	Epilog_Thunk m_EpilogThunk;
	HookSettingEntry m_HookEntry;

	void Initialize(HookSettingEntry * pHookSettingEntry, BYTE bOverwriteEntryPointLength)
	{
		assert(pHookSettingEntry!=NULL);
		if(pHookSettingEntry==NULL)
		{
			ArgumentException ExceptObj("HookRecordList::Initialize(): pHookSettingEntry is NULL");
			throw ExceptObj;
		}
		BYTE * pAPIAddress = (BYTE * )pHookSettingEntry->m_dwpFunctionAddress;
		this->m_PrologThunk.SetSavedAPIEntryPointLength( bOverwriteEntryPointLength );

		m_HookEntry = *pHookSettingEntry;
 
		// The API ReEntry address is the 
		// next instruction after the  
		// overwritten API Entry
		DWORD_PTR dwAPIReEntryAddr = (DWORD_PTR)(pAPIAddress + bOverwriteEntryPointLength);

		DWORD_PTR dwNextAddressOfJmpInstruction = this->m_PrologThunk.GetNextAddressOfJmpThunkToAPIReEntry();
		this->m_PrologThunk.SetAPIReEntryRelativeAddress( dwAPIReEntryAddr, dwNextAddressOfJmpInstruction );
		this->m_PrologThunk.SaveAPIEntryPoint( pAPIAddress, bOverwriteEntryPointLength );
	}
};
#pragma pack()


class SharedStateTLS
{
public:
	DWORD s_dwHookedAPIRecursiveStackTlsIndex;

	SharedStateTLS()
	{
		s_dwHookedAPIRecursiveStackTlsIndex = TlsAlloc();
		assert(s_dwHookedAPIRecursiveStackTlsIndex!=TLS_OUT_OF_INDEXES);
		if(s_dwHookedAPIRecursiveStackTlsIndex==TLS_OUT_OF_INDEXES)
		{
			AllocationFailureException ExceptionObj("TlsAlloc fails with error %d\n",
				GetLastError());
			throw ExceptionObj;
		}
	}
	~SharedStateTLS()
	{
		BOOL fTlsFree = TlsFree(s_dwHookedAPIRecursiveStackTlsIndex);
		assert(fTlsFree);
		if( !fTlsFree )
		{
			Win32APIFailureException ExceptionObj("TlsFree fails with error %d\n",
				GetLastError());
			throw ExceptionObj;
		}
	}
};

class HookedAPIRecursiveStack
{
private:
	const static DWORD MaxStackRecursiveLevel = 100;
	DWORD m_dwStackEntryCount;
	HookedAPIStackFrame m_stackFrameArray[MaxStackRecursiveLevel];
	CRITICAL_SECTION m_csHookedAPIRecursiveStack;
	//DWORD_PTR m_dwReturnAddressArray[MaxStackRecursiveLevel];
	//HookSettingEntry * m_HookSettingEntryArray[MaxStackRecursiveLevel];
	
	static SharedStateTLS m_ObjSharedStateTLS;
public:
	HookedAPIRecursiveStack(): m_dwStackEntryCount(0)
	{
		::InitializeCriticalSection(&m_csHookedAPIRecursiveStack);
	}

	~HookedAPIRecursiveStack()
	{
		::DeleteCriticalSection(&m_csHookedAPIRecursiveStack);
	}

	void Push(HookedAPIStackFrame objStackFrame)
	{
		assert( this->m_dwStackEntryCount < MaxStackRecursiveLevel );
		if(this->m_dwStackEntryCount>=MaxStackRecursiveLevel)
		{
			FatalException ExceptionObj("HookedAPIRecursiveStack is full!\n");
			throw ExceptionObj;
		}

		::EnterCriticalSection(&m_csHookedAPIRecursiveStack);
		this->m_stackFrameArray[this->m_dwStackEntryCount] = objStackFrame;
		this->m_dwStackEntryCount++;
		::LeaveCriticalSection(&m_csHookedAPIRecursiveStack);
	}

	HookedAPIStackFrame Pop()
	{
		assert( this->m_dwStackEntryCount>0 );
		if(this->m_dwStackEntryCount<=0)
		{
			FatalException ExceptionObj("HookedAPIRecursiveStack is full!\n");
			throw ExceptionObj;
		}

		::EnterCriticalSection(&m_csHookedAPIRecursiveStack);
		this->m_dwStackEntryCount--;
		HookedAPIStackFrame stackFrame = this->m_stackFrameArray[this->m_dwStackEntryCount];
		::LeaveCriticalSection(&m_csHookedAPIRecursiveStack);

		return stackFrame;
	}

	static HookedAPIRecursiveStack * GetInstance()
	{
		//Get the per-thread API call return address stack
		//If this is the first time, we should allocate the stack memory first
		HookedAPIRecursiveStack * pHookedAPIRecursiveStack = NULL;

		if( TlsGetValue( m_ObjSharedStateTLS.s_dwHookedAPIRecursiveStackTlsIndex )==NULL )
		{
			pHookedAPIRecursiveStack = new HookedAPIRecursiveStack();
			assert( pHookedAPIRecursiveStack );
			
			TlsSetValue(m_ObjSharedStateTLS.s_dwHookedAPIRecursiveStackTlsIndex, pHookedAPIRecursiveStack);
		}
		pHookedAPIRecursiveStack = (HookedAPIRecursiveStack *)TlsGetValue( m_ObjSharedStateTLS.s_dwHookedAPIRecursiveStackTlsIndex );

		return pHookedAPIRecursiveStack;
	}

	static void Release()
	{
		LPVOID pTlsSlot = TlsGetValue( m_ObjSharedStateTLS.s_dwHookedAPIRecursiveStackTlsIndex );
		if(pTlsSlot)
		{
			HookedAPIRecursiveStack * pHookedAPIRecursiveStack = (HookedAPIRecursiveStack *)pTlsSlot;
			delete pHookedAPIRecursiveStack;
		}
	}
};

void FreeTLSMemory()
{
	HookedAPIRecursiveStack::Release();
}

SharedStateTLS HookedAPIRecursiveStack::m_ObjSharedStateTLS;
static HookRecordList * s_pHookRecordList;


// Keeping a list of all hooked APIs by this hook engine
// Each hooked API with one corresponding HookRecord in the list
class HookRecordList
{
private:
	static HANDLE s_hHeap;
public:
	HookRecordList():dwHookRecordCount(0)
	{
		InitializeCriticalSection(&m_csHookRecordList);
		memset(dwHookRecordKeyArray, 0, sizeof(dwHookRecordKeyArray));
	}

	~HookRecordList()
	{
		DeleteCriticalSection(&m_csHookRecordList);
	}

	void * operator new(size_t mem_size, HANDLE hHeap)
	{
		assert(hHeap!=INVALID_HANDLE_VALUE);
		assert(hHeap!=NULL);

		s_hHeap = hHeap;
		LPVOID pExecutableMemory = HeapAlloc(hHeap, HEAP_ZERO_MEMORY,
			1024*X86_PAGE_SIZE);
		assert(pExecutableMemory!=NULL);
		if(pExecutableMemory==NULL)
		{
			TRACE(TEXT("HeapAlloc fails with error %d\n"),
				GetLastError());
			throw std::bad_alloc();
		}

		return pExecutableMemory;
	}

	void operator delete(void * pMemory, HANDLE hHeap)
	{
		if(!pMemory)
		{
			return;
		}

		BOOL fHeapFree = HeapFree(hHeap, 0, pMemory);
		assert(fHeapFree);
	}

	void operator delete(void * pMemory)
	{
		if(!pMemory)
		{
			return;
		}

		BOOL fHeapFree = HeapFree(s_hHeap, 0, pMemory);
		assert(fHeapFree);
	}

	BOOL InitializeNewHookRecord(HookSettingEntry * pHookSettingEntry, DWORD_PTR dwpHookRecordKey, BYTE bOverwriteEntryPointLength)
	{
		//!!!!!!!!!!!!!!!!!!
		//Warning: we should perform protection for the shared detours stack
		//!!!!!!!!!!!!!!!!!!
		::EnterCriticalSection(&m_csHookRecordList);

		this->dwHookRecordCount++;
		assert( dwHookRecordCount<=MAX_HOOK_RECORD_COUNT );
		if( dwHookRecordCount>MAX_HOOK_RECORD_COUNT )
		{
			FatalException ExceptionObj("The HookRecordList is full! The max number is MAX_HOOK_RECORD_COUNT(100)\n");
			throw ExceptionObj;
		}

		DWORD dwNextEmptyHookRecordIndex, dwIndex = 0;
		while(this->dwHookRecordKeyArray[dwIndex]!=0)
		{
			dwIndex++;
		}
		dwNextEmptyHookRecordIndex = dwIndex;
		assert( dwNextEmptyHookRecordIndex<MAX_HOOK_RECORD_COUNT );

		this->dwHookRecordKeyArray[dwNextEmptyHookRecordIndex] = dwpHookRecordKey;

		HookRecord * pHookRecord = &(this->HookRecordArray[dwNextEmptyHookRecordIndex]);
		pHookRecord->Initialize(pHookSettingEntry,bOverwriteEntryPointLength);

		::LeaveCriticalSection(&m_csHookRecordList);
		return TRUE;
	}

	BOOL FreeHookRecord(DWORD_PTR dwpHookRecordKey)
	{
		::EnterCriticalSection(&m_csHookRecordList);
		DWORD dwHookRecordIndex;
		HookRecord * pHookRecord = this->FindHookRecord(dwpHookRecordKey, &dwHookRecordIndex);
		assert( pHookRecord!=NULL );

		this->dwHookRecordCount--;
		this->dwHookRecordKeyArray[dwHookRecordIndex] = 0; 
		::LeaveCriticalSection(&m_csHookRecordList);

		return TRUE;
	}

	HookRecord * FindHookRecord(DWORD_PTR dwpHookRecordKey)
	{
		HookRecord * pHookRecord = this->FindHookRecord(dwpHookRecordKey, NULL);
		return pHookRecord;
	}
private:
	DWORD dwHookRecordCount;
	DWORD_PTR dwHookRecordKeyArray[MAX_HOOK_RECORD_COUNT];
	HookRecord HookRecordArray[MAX_HOOK_RECORD_COUNT];
	CRITICAL_SECTION m_csHookRecordList;

	HookRecord * FindHookRecord(DWORD_PTR dwpHookRecordKey, PDWORD pdwIndex)
	{
		//Enumerate the detour stack to find the detour buffer using dwpHookRecordKey
		DWORD dwIndex = 0, dwCount = 0;
		while(dwCount<this->dwHookRecordCount)
		{
			if(this->dwHookRecordKeyArray[dwIndex]!=0)
			{
				if(this->dwHookRecordKeyArray[dwIndex] == dwpHookRecordKey)
				{
					if(pdwIndex!=NULL)
					{
						*pdwIndex = dwIndex;
					}
					return &(this->HookRecordArray[dwIndex]);
				}
				dwCount++;
			}
			dwIndex++;
		}
		if(pdwIndex!=NULL)
		{
			*pdwIndex = (DWORD)-1;
		}
		return NULL;
	}
};
HANDLE HookRecordList::s_hHeap=NULL;


//=============================================================================
// Entering:
//		pHookedAPIReturnAddress[0]: Epilog()'s return address
//		pHookedAPIReturnAddress[1]: The address of the call thunk to ProxyEpilog
// 
// Exiting:
//		pHookedAPIReturnAddress[0]: Epilog()'s return address
//		pHookedAPIReturnAddress[1]: Real return address of the hooked API
// ============================================================================
//	Work:
//		0. We will not tough the pHookedAPIReturnAddress[0]
//		1. Restore the API real return address for the hooked API
//		2. Logging Exiting
// ============================================================================
void __stdcall Epilog(DWORD_PTR * pdwHookedAPIReturnAddress)
{
	HookedAPIRecursiveStack * pHookedAPIRecursiveStack = HookedAPIRecursiveStack::GetInstance();

	HookedAPIStackFrame objStackFrame = pHookedAPIRecursiveStack->Pop();
	pdwHookedAPIReturnAddress[1] = objStackFrame.m_dwpReturnAddress;

	TRACE(TEXT("TID: %d, [%s!%s] at address 0x%p exits\n"), 
		GetCurrentThreadId(), 
		objStackFrame.m_pHookRecord->m_HookEntry.m_szModuleName, 
		objStackFrame.m_pHookRecord->m_HookEntry.m_szFunctionName, 
		objStackFrame.m_pHookRecord->m_HookEntry.m_dwpFunctionAddress);

	DWORD_PTR dwpHookRecordKey = objStackFrame.m_pHookRecord->m_HookEntry.m_dwpFunctionAddress+SHORT_CALL_INSTRUCTION_LENGTH;
	ReferenceCountEntry * pCountEntry = s_pHookAPIReferenceCountList->FindEntry(dwpHookRecordKey);
	assert(pCountEntry!=NULL);
	if(pCountEntry==NULL)
	{
		FatalException ExceptObj("Epilog(): can't find ReferenceCountEntry\n");
		throw ExceptObj;
	}

	pCountEntry->Decrease();
}


const DWORD bOffsetToSavedEax = 12;
__declspec(naked) void ProxyEpilog()
{
	__asm
	{
		push eax
		push ebx
		push ecx
		push edx

		mov ebx, esp
		pushfd

		add ebx, bOffsetToSavedEax
		push ebx


		//Note: Epilog is the __stdcall, so the parameter will be clean
		//up by the Epilog
		Call Epilog
		
		popfd
		pop edx
		pop ecx
		pop ebx
		pop eax
		ret
	}
}


//=============================================================================
// Entering:
//		pHookedAPIReturnAddress[0]: The key to retrieve Hook_Record
//		pHookedAPIReturnAddress[1]: Real return address of the hooked API
// 
// Exiting:
//		pHookedAPIReturnAddress[0]: The address of the saved API EntryPoint Buffer
//		pHookedAPIReturnAddress[1]: The address of the call thunk to ProxyEpilog
//										which allows the ProxyEpilog to get control
//										during API exiting.
// ============================================================================
//	The execution will jmp to pHookedAPIReturnAddress[0] after ProxyProlog
// ============================================================================
void __stdcall Prolog(DWORD_PTR * pHookedAPIReturnAddress)
{
	DWORD_PTR dwpHookRecordKey = pHookedAPIReturnAddress[0];

	ReferenceCountEntry * pCountEntry = s_pHookAPIReferenceCountList->FindEntry(dwpHookRecordKey);
	assert(pCountEntry!=NULL);
	if(pCountEntry==NULL)
	{
		FatalException ExceptObj("Prolog(): can't find ReferenceCountEntry\n");
		throw ExceptObj;
	}

	DWORD dwNewCount = pCountEntry->CheckIncrease();
	if(dwNewCount>0)
	{
		HookRecord * pHookRecord = s_pHookRecordList->FindHookRecord(dwpHookRecordKey);
		assert(pHookRecord!=NULL);
		if(pHookRecord==NULL)
		{
			FatalException ExceptObj("Prolog(): can't find HookRecord\n");
			throw ExceptObj;
		}

		HookedAPIRecursiveStack * pHookedAPIRecursiveStack = HookedAPIRecursiveStack::GetInstance();
		assert( pHookedAPIRecursiveStack!=NULL );
		if(pHookedAPIRecursiveStack==NULL)
		{
			OutOfMemoryException ExceptObj("Prolog(): constructing HookedAPIRecursiveStack() failed\n");
			throw ExceptObj;
		}

		HookedAPIStackFrame objStackFrame(pHookedAPIReturnAddress[1], pHookRecord);
		pHookedAPIRecursiveStack->Push(objStackFrame);

		pHookedAPIReturnAddress[0] = (DWORD_PTR)(pHookRecord->m_PrologThunk.GetSavedAPIEntryPointBufferAddress());
		pHookedAPIReturnAddress[1] = (DWORD_PTR)(pHookRecord->m_EpilogThunk.GetCallThunkToProxyEpilogAddress());


		TRACE(TEXT("TID: %d, [%s!%s] at address 0x%p enters\n"), 
			GetCurrentThreadId(), 
			pHookRecord->m_HookEntry.m_szModuleName, 
			pHookRecord->m_HookEntry.m_szFunctionName, 
			pHookRecord->m_HookEntry.m_dwpFunctionAddress);

		const DWORD dwFramesToCapture = 65;
		DWORD_PTR StackFrames[dwFramesToCapture];
		DWORD dwFramesCatpured = CaptureStackBackTrace(0, dwFramesToCapture, (PVOID *)StackFrames, NULL);
		assert(dwFramesCatpured>0);

		StackTrace * pCurrentStackTrace = AllocateNewStackTrace(dwFramesCatpured+1);
		assert(pCurrentStackTrace);
		if(pCurrentStackTrace==NULL)
		{
			FatalException ExceptObj("Prolog(): AllocateNewStackTrace() returns NULL\n");
			throw ExceptObj;
		}

		pCurrentStackTrace->m_dwpStackFrames[0] = pHookRecord->m_HookEntry.m_dwpFunctionAddress;
		for(DWORD dwIndex=0;dwIndex<dwFramesCatpured;dwIndex++)
		{
			pCurrentStackTrace->m_dwpStackFrames[dwIndex+1] = StackFrames[dwIndex];
		}

		//TRACEA(szStackTraceBuffer);*/
	}
	else
	{
		pHookedAPIReturnAddress[0] = dwpHookRecordKey-SHORT_CALL_INSTRUCTION_LENGTH;
	}
}


// Since call thunk(which overwrote the API Entry 
// the top of the stack contains:
//		1. Next instruction address after call thunk
//		2. Final return address of the hooked API
// These two values will be passed as parameters
// to the "Prolog"
const DWORD bSavedGeneralRegisterLength = 16;
__declspec(naked) void ProxyProlog()
{
	__asm
	{
		push eax
		push ebx
		push ecx
		push edx

		// ebx points to the top of the stack.
		// (Note: currently, the Eflags is not saved yet)
		mov ebx, esp
		pushfd


		// Ebx jmps over the saved Registers and points to 
		// the top of the stack before saving registers
		add ebx, bSavedGeneralRegisterLength

		// This pointer is the parameter to Prolog
		push ebx

		// Note: Prolog is the __stdcall, 
		// we do not need to clean the parameters
		Call Prolog

		popfd
		pop edx
		pop ecx
		pop ebx
		pop eax


		// This will transfer the execution to
		// the top of the stack which is 
		// SavedAPIEntryAddress(modified in Prolog)
		ret
	}
}


DWORD HookEngine::SetMemoryBlockProtection(LPCVOID pMemoryAddress, DWORD dwNewProtection)
{
	MEMORY_BASIC_INFORMATION mbi;

	SIZE_T uVirtualQueryResult = VirtualQuery(pMemoryAddress, &mbi, sizeof(mbi));
	assert( uVirtualQueryResult== sizeof(mbi) );
	if( uVirtualQueryResult!= sizeof(mbi))
	{
		Win32APIFailureException ExceptionObj("VirtualQuery fails with error %d\n",
			GetLastError());
		throw ExceptionObj;
	}

	DWORD dwOldProtection;
	BOOL fVirtualProtectResult = VirtualProtect(mbi.BaseAddress, mbi.RegionSize,
		dwNewProtection, &dwOldProtection);
	assert( fVirtualProtectResult );
	if(!fVirtualProtectResult)
	{
		Win32APIFailureException ExceptionObj("VirtualProtect fails with error %d\n",
			GetLastError());
		throw ExceptionObj;
	}

	return dwOldProtection;
}

BOOL HookEngine::OverwriteAPIEntryPoint(LPVOID pAPIEntryAddress, LPCVOID pOverwrittenInstructionBuf, SIZE_T nOverwrittenInstructionBufLength )
{
	DWORD dwBytesWritten;
	BOOL fWriteProcessMemoryResult = WriteProcessMemory(GetCurrentProcess(), pAPIEntryAddress, pOverwrittenInstructionBuf, 
		nOverwrittenInstructionBufLength, &dwBytesWritten);
	assert( fWriteProcessMemoryResult );
	if(!fWriteProcessMemoryResult)
	{
		Win32APIFailureException ExceptionObj("WriteProcessMemory fails with error %d\n",
			GetLastError());
		throw ExceptionObj;
	}

	return TRUE;
}

class HookEngineThread
{
private:
	HANDLE m_hThread;
public:
	HookEngineThread(DWORD dwThreadId)
	{
		m_hThread = OpenThread(THREAD_SUSPEND_RESUME|THREAD_GET_CONTEXT, FALSE, dwThreadId);
		//assert(m_hThread!=NULL);
		if(m_hThread==NULL)
		{
			Win32APIFailureException peObj("OpenThread(THREAD_SUSPEND_RESUME|THREAD_GET_CONTEXT) fails with error %d\n", GetLastError());
			throw peObj;
		}
	}

	DWORD Suspend()
	{
		DWORD dwSuspendCount = SuspendThread(m_hThread);
		assert(dwSuspendCount!=(DWORD)-1);
		return dwSuspendCount;
	}

	DWORD Resume()
	{
		DWORD dwSuspendCount = ResumeThread(m_hThread);
		assert(dwSuspendCount!=(DWORD)-1);
		return dwSuspendCount;
	}


	DWORD_PTR GetEip()
	{
		CONTEXT threadContext;
		threadContext.ContextFlags = CONTEXT_CONTROL;
		BOOL fResult = GetThreadContext(m_hThread, &threadContext);
		assert(fResult);
		if(!fResult)
		{
			return 0;
		}
	
		return threadContext.Eip;	
	}

	BOOL SetEip(DWORD_PTR dwpEip)
	{
		CONTEXT threadContext;
		threadContext.ContextFlags = CONTEXT_CONTROL;
		threadContext.Eip = dwpEip;

		BOOL fResult = SetThreadContext(m_hThread, &threadContext);
		assert(fResult);
		if(!fResult)
		{
			return FALSE;
		}
	
		return TRUE;	
	}

	~HookEngineThread()
	{
		CloseHandle(m_hThread);
	}
};

BOOL HookEngine::SuspendAllThreads()
{

	HANDLE hThreadSnapShot = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
	assert(INVALID_HANDLE_VALUE!=hThreadSnapShot);
	if(INVALID_HANDLE_VALUE==hThreadSnapShot)
	{
		Win32APIFailureException ExceptionObj("CreateToolhelp32Snapshot fails with error %d\n",
			GetLastError());
		throw ExceptionObj;
	}

	THREADENTRY32 objThread32 ={ sizeof(THREADENTRY32)};
	if(!Thread32First(hThreadSnapShot, &objThread32))
	{
		DWORD dwErrorCode = GetLastError();
		CloseHandle(hThreadSnapShot);

		Win32APIFailureException ExceptionObj("Thread32First fails with error %d\n",
			dwErrorCode);
		throw ExceptionObj;
	}

	DWORD dwCurrentProcessId = GetCurrentProcessId();
	do
	{
		if(objThread32.th32OwnerProcessID == dwCurrentProcessId)
		{
			if(objThread32.th32ThreadID!=GetCurrentThreadId())
			{
				HookEngineThread objHookThread(objThread32.th32ThreadID);
				objHookThread.Suspend();
			}
		}
	}while(Thread32Next(hThreadSnapShot, &objThread32));

	return TRUE;
}

BOOL HookEngine::ResumeAllThreads()
{
	HANDLE hThreadSnapShot = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
	assert(INVALID_HANDLE_VALUE!=hThreadSnapShot);
	if(INVALID_HANDLE_VALUE==hThreadSnapShot)
	{
		Win32APIFailureException ExceptionObj("CreateToolhelp32Snapshot fails with error %d\n",
			GetLastError());
		throw ExceptionObj;
	}

	THREADENTRY32 objThread32 ={ sizeof(THREADENTRY32)};
	if(!Thread32First(hThreadSnapShot, &objThread32))
	{
		DWORD dwErrorCode = GetLastError();
		CloseHandle(hThreadSnapShot);

		Win32APIFailureException ExceptionObj("Thread32First fails with error %d\n",
			dwErrorCode);
		throw ExceptionObj;
	}

	BOOL fResult = FlushInstructionCache(GetCurrentProcess(), NULL, NULL);
	assert(fResult);

	DWORD dwCurrentProcessId = GetCurrentProcessId();
	do
	{
		if(objThread32.th32OwnerProcessID == dwCurrentProcessId)
		{
			if(objThread32.th32ThreadID!=GetCurrentThreadId())
			{
				HookEngineThread objHookThread(objThread32.th32ThreadID);
				objHookThread.Resume();		
			}
		}
	}while(Thread32Next(hThreadSnapShot, &objThread32));

	return TRUE;
}


BOOL HookEngine::IsEipInCodeRange(DWORD_PTR eip, DWORD_PTR dwpStartAddress, DWORD dwpCodeRangeLength)
{
	return (eip>=dwpStartAddress)&&eip<=(dwpStartAddress+dwpCodeRangeLength);
}


BOOL HookEngine::IsHookInUse(DWORD_PTR dwpAPIAddress)
{
	DWORD_PTR dwpHookRecordKey = dwpAPIAddress + SHORT_CALL_INSTRUCTION_LENGTH;
	if(s_pHookAPIReferenceCountList->FindEntry(dwpHookRecordKey)->GetCount()>1)
		return TRUE;

	/*HANDLE hThreadSnapShot = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, GetCurrentProcessId());
	assert(INVALID_HANDLE_VALUE!=hThreadSnapShot);
	if(INVALID_HANDLE_VALUE==hThreadSnapShot)
	{
		Win32APIFailureException exceptionObj("CreateToolhelp32Snapshot fails with error %d\n",
				GetLastError());
		throw exceptionObj;
	}

	THREADENTRY32 objThread32 ={ sizeof(THREADENTRY32)};
	if(!Thread32First(hThreadSnapShot, &objThread32))
	{
		CloseHandle(hThreadSnapShot);
		Win32APIFailureException exceptionObj("Thread32First fails with error %d\n",
				GetLastError());
		throw exceptionObj;
	}

	DWORD dwCurrentProcessId = GetCurrentProcessId();
	do
	{
		if(objThread32.th32OwnerProcessID == dwCurrentProcessId)
		{
			if(objThread32.th32ThreadID!=GetCurrentThreadId())
			{
				try
				{
					HookEngineThread objHookThread(objThread32.th32ThreadID);
					DWORD_PTR dwpEip = objHookThread.GetEip();

					DWORD_PTR dwpHookRecordKey = dwpAPIAddress + SHORT_CALL_INSTRUCTION_LENGTH;
					assert(s_pHookRecordList!=NULL);
					HookRecord * pHookRecord = s_pHookRecordList->FindHookRecord(dwpHookRecordKey);
					assert(pHookRecord!=NULL);

					HINSTANCE hHookEngine = GetModuleHandle(TEXT("HookEngine.dll"));
					assert(hHookEngine!=NULL);

					MODULEINFO mi;
					BOOL fResult = GetModuleInformation(GetCurrentProcess(), hHookEngine, &mi, sizeof(MODULEINFO));
					assert(fResult);

					if(this->IsEipInCodeRange(dwpEip, (DWORD_PTR)hHookEngine, mi.SizeOfImage))
						return TRUE;
				}
				catch(Win32APIFailureException&)
				{}
			}
		}
	}while(Thread32Next(hThreadSnapShot, &objThread32));*/

	return FALSE;
}

BOOL HookEngine::UpdateAllThreads(DWORD_PTR dwpAPIAddress, HookOperation operation)
{
	HANDLE hThreadSnapShot = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, GetCurrentProcessId());
	assert(INVALID_HANDLE_VALUE!=hThreadSnapShot);
	if(INVALID_HANDLE_VALUE==hThreadSnapShot)
	{
		Win32APIFailureException ExceptionObj("CreateToolhelp32Snapshot fails with error %d\n",
			GetLastError());
		throw ExceptionObj;
	}

	THREADENTRY32 objThread32 ={ sizeof(THREADENTRY32)};
	if(!Thread32First(hThreadSnapShot, &objThread32))
	{
		DWORD dwErrorCode = GetLastError();
		CloseHandle(hThreadSnapShot);

		Win32APIFailureException ExceptionObj("Thread32First fails with error %d\n",
			dwErrorCode);
		throw ExceptionObj;
	}

	DWORD dwCurrentProcessId = GetCurrentProcessId();
	do
	{
		if(objThread32.th32OwnerProcessID == dwCurrentProcessId)
		{
			if(objThread32.th32ThreadID!=GetCurrentThreadId())
			{
				HookEngineThread objHookThread(objThread32.th32ThreadID);
				DWORD_PTR dwpEip = objHookThread.GetEip();

				DWORD_PTR dwpHookRecordKey = dwpAPIAddress + SHORT_CALL_INSTRUCTION_LENGTH;

				assert(s_pHookRecordList!=NULL);
				if(s_pHookRecordList==NULL)
				{
					FatalException ExceptObj("HookEngine::UpdateAllThreads: s_pHookRecordList is NULL\n");
					throw ExceptObj;
				}
				HookRecord * pHookRecord = s_pHookRecordList->FindHookRecord(dwpHookRecordKey);
				assert(pHookRecord!=NULL);
				if(pHookRecord==NULL)
				{
					FatalException ExceptObj("HookEngine::UpdateAllThreads: can't find HookRecord(pHookRecord is NULL)\n");
					throw ExceptObj;
				}

				DWORD dwSavedAPIEntryPointLength = pHookRecord->m_PrologThunk.GetSavedAPIEntryPointLength();

				DWORD_PTR dwpNewEip = 0;
				if(operation == HookAdd )
				{
					if(this->IsEipInCodeRange(dwpEip , dwpAPIAddress, dwSavedAPIEntryPointLength))
					{ 
						dwpNewEip = pHookRecord->m_PrologThunk.GetSavedAPIEntryPointBufferAddress() 
							+ (dwpEip - dwpAPIAddress);
						
						assert(dwpNewEip!=0);
						objHookThread.SetEip(dwpNewEip);
					}
				}
				else if(operation == HookRemove)
				{
					/*DWORD_PTR dwpSavedAPIEntryPointBufferStartAddress = pHookRecord->m_PrologThunk.GetSavedAPIEntryPointBufferAddress();
					DWORD_PTR dwpBufferLength = pHookRecord->m_PrologThunk.GetSavedAPIEntryPointBufferLength();

					if(this->IsEipInCodeRange(dwpEip, 
						dwpSavedAPIEntryPointBufferStartAddress, 
						dwSavedAPIEntryPointLength
						))
					{
						dwpNewEip = dwpAPIAddress
							+ (dwpEip - dwpSavedAPIEntryPointBufferStartAddress);

						assert(dwpNewEip!=0);
						objHookThread.SetEip(dwpNewEip);
					}
					else if(this->IsEipInCodeRange(dwpEip, 
						dwpSavedAPIEntryPointBufferStartAddress,
						dwpBufferLength
						))
					{
						 dwpNewEip = dwpAPIAddress + dwSavedAPIEntryPointLength;
						 
						 assert(dwpNewEip!=0);
						 objHookThread.SetEip(dwpNewEip);
					}*/
				}
			}
		}
	}while(Thread32Next(hThreadSnapShot, &objThread32));

	return TRUE;
}


// Disassembly the target API to check its prolog length
// Also check if the API is already hooked
BYTE HookEngine::GetOverwrittenAPIEntryPointLength(BYTE * pAPIAddr)
{
	BYTE bOverwriteEntryPointLength = 0;
	INSTRUCTION inst;

	while(bOverwriteEntryPointLength < SHORT_CALL_INSTRUCTION_LENGTH)
	{
		int len = get_instruction(&inst, pAPIAddr+bOverwriteEntryPointLength, MODE_32);
		assert( len>0 );
		assert( inst.type != INSTRUCTION_TYPE_JMP );
		assert( inst.type != INSTRUCTION_TYPE_CALL );
		assert( inst.type != INSTRUCTION_TYPE_RET );

		if( len<=0 )
		{
			FatalException ExceptionObj("get_instruction fails!\n");
			throw ExceptionObj;
		}

		if((inst.type == INSTRUCTION_TYPE_JMP)||inst.type == INSTRUCTION_TYPE_CALL)
		{
			TRACE(TEXT("There is already hook applied at address 0x%08X\n"),
				pAPIAddr);
			return 0;
		}
		else if(inst.type == INSTRUCTION_TYPE_RET)
		{
			char szErrMsg[MaxErrMsgLen];
			sprintf_s(szErrMsg, sizeof(szErrMsg), "Function length less than SHORT_CALL_INSTRUCTION_LENGTH(%d bytes)\n", SHORT_CALL_INSTRUCTION_LENGTH);
			
			FatalException ExceptionObj(szErrMsg);
			throw ExceptionObj;
		}
		bOverwriteEntryPointLength += len;
	}

	// Warning: bOverwriteEntryPointLength!=SHORT_CALL_INSTRUCTION_LENGTH
	// bOverwriteEntryPointLength must count for the instruction boundary(entire instruction 
	// length that is overwritten).
	return bOverwriteEntryPointLength;
}


BOOL HookEngine::Start(HookSettingEntry * pHookSettingEntry)
{
	assert( pHookSettingEntry!= NULL);
	if(pHookSettingEntry==NULL)
	{
		ArgumentException aeobj("HookEngine::Start: pHookSettingEntry is NULL\n");
		throw aeobj;
	}
	BYTE * pAPIAddr = (BYTE *)pHookSettingEntry->m_dwpFunctionAddress;
	assert( !IsBadCodePtr((FARPROC)pHookSettingEntry) );

	BYTE bOverwriteEntryPointLength = GetOverwrittenAPIEntryPointLength(pAPIAddr);
	assert( bOverwriteEntryPointLength>0 );
	if( bOverwriteEntryPointLength==0 )
	{
		TRACE(TEXT("GetOverwrittenAPIEntryPointLength() fails\n"));
		return FALSE;
	}

	
	// We use the "API Entry address + SHORT_CALL_INSTRUCTION_LENGTH" 
	// as the dwpHookRecordKey to unique identify the HookRecord 
	// This will help us to use this key 
	// in the "Prolog" to retrieve the HookRecord
	DWORD_PTR dwNextAddressOfAPIEntryPointCallThunk = (DWORD_PTR)pAPIAddr + SHORT_CALL_INSTRUCTION_LENGTH;
	DWORD_PTR dwpHookRecordKey = dwNextAddressOfAPIEntryPointCallThunk;

	BOOL fInitializeNewHookRecordResult = s_pHookRecordList->InitializeNewHookRecord( pHookSettingEntry, dwpHookRecordKey, bOverwriteEntryPointLength );
	assert( fInitializeNewHookRecordResult );
	if( !fInitializeNewHookRecordResult )
	{
		TRACE(TEXT("InitializeNewHookRecord() fails\n"));
		return FALSE;
	}

	s_pHookAPIReferenceCountList->InitializeNewEntry(dwpHookRecordKey);


	// We use a call in struction thunk
	// to overwrite the API entry 
	// which jmps to the ProxyProlog
	x86CodeThunk objAPIEntryInlineOverwrittenCallThunk(X86_SHORTCALL_OPCODE);
	objAPIEntryInlineOverwrittenCallThunk.SetOprandRelativeAddress( (DWORD_PTR)ProxyProlog, dwNextAddressOfAPIEntryPointCallThunk );

	DWORD dwOldCodeProtect = SetMemoryBlockProtection(pAPIAddr, PAGE_EXECUTE_READWRITE);
	OverwriteAPIEntryPoint(pAPIAddr, objAPIEntryInlineOverwrittenCallThunk.GetCodeThunkAddress(), 
		objAPIEntryInlineOverwrittenCallThunk.GetCodeThunkLength() );
	SetMemoryBlockProtection(pAPIAddr, dwOldCodeProtect);
	this->UpdateAllThreads((DWORD_PTR)pAPIAddr, HookAdd);

	return TRUE;
}



BOOL HookEngine::Stop(HookSettingEntry * pHookSettingEntry)
{
	assert(pHookSettingEntry!=NULL);
	if(pHookSettingEntry==NULL)
	{
		ArgumentException ExceptionObj("HookEngine::Stop: pHookSettingEntry is NULL\n");
		throw ExceptionObj;
	}
	BYTE * pAPIAddr = (BYTE *)pHookSettingEntry->m_dwpFunctionAddress;

	assert( s_pHookRecordList!=NULL );
	if(s_pHookRecordList==NULL)
	{
		FatalException ExceptionObj("The inline hook is not initialized\n");
		throw ExceptionObj;
	}

	INSTRUCTION inst;
	int len = get_instruction(&inst, pAPIAddr, MODE_32);
	assert(inst.type == INSTRUCTION_TYPE_CALL);
	if(inst.type != INSTRUCTION_TYPE_CALL)
	{
		TRACE(TEXT("There is not inline hook installed at address 0x%08X\n"),
			pAPIAddr);
		return FALSE;
	}


	if(this->IsHookInUse((DWORD_PTR)pAPIAddr))
	{
		TRACE(TEXT("The HookEngine is in use, we can not stop it\n"));
		return FALSE;
	}
	else
	{
		DWORD_PTR dwpHookRecordKey = (DWORD_PTR)pAPIAddr + SHORT_CALL_INSTRUCTION_LENGTH;
		HookRecord * pHookRecord = s_pHookRecordList->FindHookRecord(dwpHookRecordKey);
		assert( pHookRecord!=NULL );
		if(pHookRecord==NULL)
		{
			char szErrMsg[MaxErrMsgLen];
			sprintf_s(szErrMsg, sizeof(szErrMsg), "Can not find detour_buf with detour key 0x%X\n", dwpHookRecordKey);
			
			FatalException ExceptionObj(szErrMsg);
			throw ExceptionObj;
		}

		DWORD dwOldCodeProtect = SetMemoryBlockProtection( pAPIAddr, PAGE_EXECUTE_READWRITE );
		OverwriteAPIEntryPoint( pAPIAddr, (LPCVOID)pHookRecord->m_PrologThunk.GetSavedAPIEntryPointBufferAddress(), 
			pHookRecord->m_PrologThunk.GetSavedAPIEntryPointLength() );
		SetMemoryBlockProtection( pAPIAddr, dwOldCodeProtect ); 

		s_pHookAPIReferenceCountList->FindEntry(dwpHookRecordKey)->Decrease();
		s_pHookRecordList->FreeHookRecord(dwpHookRecordKey);
	}

	return TRUE;
}

class HeapCreateWrapper
{
public:
	HeapCreateWrapper():m_hHeap(NULL)
	{
		m_hHeap = HeapCreate(HEAP_CREATE_ENABLE_EXECUTE, 0, 0);
		assert(m_hHeap!=NULL);
		if(m_hHeap==NULL)
		{
			Win32APIFailureException ExceptionObj("HeapCreate fails with error %d\n",
				GetLastError());
			throw ExceptionObj;
		}
	}

	~HeapCreateWrapper()
	{
		BOOL fHeapDestroy = HeapDestroy(m_hHeap);
		assert(fHeapDestroy);
		if( !fHeapDestroy )
		{
			TRACE(TEXT("HeapDestroy fails with error %d\n"),
				GetLastError());
			return;
		}
	}

	HANDLE GetHandle()
	{
		assert(m_hHeap!=NULL);
		return m_hHeap;
	}

private:
	HANDLE m_hHeap;
};


HookEngine::HookEngine(): pHCWrapper(new HeapCreateWrapper())
{
	s_pStackTraceList = (StackTraceList *)new char[1024*X86_PAGE_SIZE];
	s_pStackTraceList->m_dwCount = 0;
	s_pCurrentStackTrace = s_pStackTraceList->m_StackTraces;

	s_pHookAPIReferenceCountList.reset(new HookAPIReferenceCountList());
	s_pHookRecordList= new (pHCWrapper->GetHandle())HookRecordList();
}

HookEngine::~HookEngine()
{
	DWORD dwStackTraceListSize = GetStackTraceListSize();
	if(dwStackTraceListSize>0)
	{
		assert(dwStackTraceListSize%4==0);

		HANDLE hProcess = GetCurrentProcess();
		assert(hProcess);

		DWORD cbNeeded, dwModuleListSize;
		HMODULE hMods[1024];

		BOOL fResult = EnumProcessModules(hProcess, hMods, sizeof(hMods), &cbNeeded);
		assert(fResult);
		if(!fResult)
		{
			Win32APIFailureException ExceptionObj("[HookEngine::~HookEngine()]--EnumProcessModules fails with error %d\n", GetLastError());
			throw ExceptionObj;
		}

		DWORD dwModulesCount = cbNeeded / sizeof(HMODULE);
		assert(dwModulesCount>0);

		dwModuleListSize = sizeof(DWORD)+dwModulesCount*sizeof(ModuleInfo);
		ModuleInfoList * pModulesList = (ModuleInfoList *)malloc(dwModuleListSize);
		assert(pModulesList);
		if(!pModulesList)
		{
			OutOfMemoryException ExceptionObj("[HookEngine::~HookEngine()]--malloc fails\n");
			throw ExceptionObj;
		}

		pModulesList->m_dwCount = dwModulesCount;

		for (DWORD dwIndex = 0; dwIndex < dwModulesCount; dwIndex++ )
		{
			pModulesList->m_modules[dwIndex].m_dwpStartAddress = (DWORD64)hMods[dwIndex];

			MODULEINFO mi;
			fResult = GetModuleInformation(hProcess, hMods[dwIndex], &mi, sizeof(MODULEINFO));
			assert(fResult);
			if(!fResult)
			{
				Win32APIFailureException ExceptionObj("[HookEngine::~HookEngine()]--GetModuleInformation fails with error %d\n", GetLastError());
				throw ExceptionObj;
			}
			pModulesList->m_modules[dwIndex].m_dwSize = mi.SizeOfImage;
			
			fResult = GetModuleFileNameExW(hProcess, hMods[dwIndex], pModulesList->m_modules[dwIndex].m_szFullPath,
									  MAX_PATH);
			assert(fResult);
			if(!fResult)
			{
				Win32APIFailureException ExceptionObj("[HookEngine::~HookEngine()]--GetModuleFileNameExA fails with error %d\n", GetLastError());
				throw ExceptionObj;
			}
		}
		ofstream stackTraceOutputFile("D:\\StackTraceOutput.txt", ios_base::out|ios_base::binary);
		assert(stackTraceOutputFile!=NULL);
		TRACEA("[HookEngine::~HookEngine()] -- ModuleInfoList size is %d\n", dwModuleListSize);
		stackTraceOutputFile.write((char *)pModulesList, dwModuleListSize);

		free(pModulesList);

		TRACEA("[HookEngine::~HookEngine()] -- dwStackTraceListSize is %d\n", dwStackTraceListSize);
		stackTraceOutputFile.write((char *)s_pStackTraceList, dwStackTraceListSize);
	}
	

	delete [] s_pStackTraceList;
	delete s_pHookRecordList;
}
